2021-ICML-[EGNN]E(n) Equivariant Graph Neural Networks

https://arxiv.org/abs/2102.09844

Introduction

データに対して、対称性を利用して(CNNによって画像が回転してもetcでとらえられるように、GNNでグラフの点を入れる順番が前後しても同じもの扱いされるように)

多くの問題では、データが3次元空間での平行移動や回転に対して対称性を持つ(分子構造や化学特性予測、物理シミュレーション、3D点群データの認識)

これらはユークリッド群 $E(3)$ の上であるといえる。

この等価性を保つには、CNNなどの手法があるがそれだけでは不十分。先行研究では中間表現に対してメスを入れてたが、その中間表現を計算するのに計算コストが高い。

この研究では、 $E(n)$ ユークリッド群にまで拡張した新しいアーキテクチャを提案した。

Background

同変性(Equivariance)

ある抽象群 $g \in G$ について考える。同変性を持つ関数 $\phi : X \to Y$ とは以下のようなもの。
- 入力空間 $X$ についての変換 $T_g : X \to X$ と、出力空間 $Y$ について、変換 $S_g : Y \to Y$ が存在する。
- 以下のように、 $T_g$ してから $\phi$ で $X \to Y$ するのと、 $\phi$ した結果 $Y$ になったのを $S_g$ でまた変換しても、同じ値となる。
$\phi(T_g(x)) = S_g(\phi(x))$
- 具体例として、平行移動同変性がある。 $g$ だけ入力を平行移動させると、 $g$ だけ出力を平行移動させないといけないという条件を、 $\phi$ は満たさないといけない。
$T_g(x) = x + g \\ S_g(x) = y + g$
機械学習でよく考えられる、同変性には3つある。
- 平行移動同変性。平行移動させても同じ。
$x + g = (x_1 + g, \cdots, x_n + g) , y + g = \phi(x + g)$
- 回転同変性。 $Q$ という直交行列で回転を実現。回転させても同じ。
$Qx = (Qx_1, \cdots, Q x_M), Qy = \phi(Qx)$
- 順列同変性。順列を変えても同じ。入力についての順列 $P(x)$ について。
$P(y) = \phi(P(x))$

Graph Neural Network

GNNは順列同変性を持つ、グラフについてのDNNである。

このように、隣接した頂点と、その辺にある情報を受け取り $m_{i,j}$ 、それを集計することで $m_i$ となり、最終的に次の自分の頂点の要素 $h_i^{l+1}$ は今の要素 $h_i^l$ と周辺の情報 $m_i$ によって作られる。 $\phi_e, \phi_h$ は一般的にMLPで実現。

同変性を持つGNN

Equivariant Graph Neural Network(EGNN)

入力として、以下のものを受け取る。

$\mathbf{h}^l = \{ \mathbf{h}_0^l, \cdots, \mathbf{h}_{M-1}^l \}$ は各ノードの埋め込みベクトル。
$\mathbf{x}^l = \{ \mathbf{x}_0^l, \cdots, \mathbf{x}_{M-1}^l \}$ は各ノードの座標埋め込みベクトル。
$\Epsilon$ はグラフの辺の情報。辺は引き続き重みもある。

辺の埋め込みに加えて、2つの座標埋め込み $\mathbf{x}_i, \mathbf{x}_j$ でユークリッド距離を計算し、それを辺の埋め込みの代わりにする。(つまり $\phi_e$ の入力は増える)
周辺ノードからの情報の集約、自分自身の更新はそのまま。
座標の更新は、新たなMLPの $\phi_x$ を用いて、そこには辺 $ij$ の更新情報埋め込み $\mathbf{m}_{ij}$ を入れる。そして、今の点 $ij$ の間の座標埋め込みの差を乗じて、最後に合算する。
- これが毎回の更新の差分。